package com.gitee.wsl.platform.network

import java.io.*
import java.net.DatagramPacket
import java.net.InetAddress
import java.net.MulticastSocket
import java.net.SocketTimeoutException


/**
 * Low-level functionality for handling questions and answers over the Multicast DNS Protocol
 * (RFC 6762), in conjunction with DNS Service Discovery (DNS-SD, RFC 6763).
 */
object MDNSDiscover {
    private const val QTYPE_A: Short = 0x0001
    const val QTYPE_PTR: Short = 0x000c
    const val QTYPE_TXT: Short = 0x0010
    const val QTYPE_SRV: Short = 0x0021

    const val QCLASS_INTERNET: Short = 0x0001
    const val CLASS_FLAG_MULTICAST: Short = 0
    val CLASS_FLAG_UNICAST: Short = 0x8000.toShort()
    private const val PORT = 5353

    private const val MULTICAST_GROUP_ADDRESS = "224.0.0.251"

    private const val DEBUG = false

    @Throws(IOException::class)
    private fun discoverPacket(serviceType: String): ByteArray {
        return queryPacket(serviceType, QCLASS_INTERNET.toInt() or CLASS_FLAG_UNICAST.toInt(), QTYPE_PTR.toInt())
    }

    @Throws(IOException::class)
    fun queryPacket(serviceName: String, qclass: Int, vararg qtypes: Int): ByteArray {
        val bos = ByteArrayOutputStream()
        val dos = DataOutputStream(bos)
        dos.writeInt(0)
        dos.writeShort(qtypes.size) // questions
        dos.writeShort(0) // answers
        dos.writeShort(0) // nscount
        dos.writeShort(0) // arcount
        var fqdnPtr = -1
        for (qtype in qtypes) {
            if (fqdnPtr == -1) {
                fqdnPtr = dos.size()
                writeFQDN(serviceName, dos)
            } else {
                // packet compression, string is just a pointer to previous occurrence
                dos.write(0xc0 or (fqdnPtr shr 8))
                dos.write(fqdnPtr and 0xFF)
            }
            dos.writeShort(qtype)
            dos.writeShort(qclass)
        }
        dos.close()
        return bos.toByteArray()
    }

    /**
     * Sends a discovery packet for the specified service and listens for reply packets, notifying
     * a callback as services are discovered.
     *
     * @param serviceType the type of service to query in mDNS, e.g. `"_example._tcp.local"`
     * @param callback    receives callbacks with [Result] objects as answers are decoded from
     * incoming reply packets.
     * @param timeout     duration in milliseconds to wait for answer packets. If `0`, this method
     * will listen forever.
     * @throws IOException
     */
    @Throws(IOException::class)
    fun discover(serviceType: String, callback: Callback?, timeout: Int) {
        require(timeout >= 0)
        val group = InetAddress.getByName(MULTICAST_GROUP_ADDRESS)
        val sock = MulticastSocket() // binds to a random free source port
        if (DEBUG) println("Source port is " + sock.localPort)
        val data = discoverPacket(serviceType)
        if (DEBUG) println("Query packet:")
        if (DEBUG) hexdump(data, 0, data.size)
        var packet = DatagramPacket(data, data.size, group, PORT)
        sock.timeToLive = 255
        sock.send(packet)
        val buf = ByteArray(1024)
        packet = DatagramPacket(buf, buf.size)
        var endTime: Long = 0
        if (timeout != 0) {
            endTime = System.currentTimeMillis() + timeout
        }
        while (true) {
            if (timeout != 0) {
                val remaining = (endTime - System.currentTimeMillis()).toInt()
                if (remaining <= 0) {
                    break
                }
                sock.soTimeout = remaining
            }
            try {
                sock.receive(packet)
            } catch (e: SocketTimeoutException) {
                break
            }
            if (DEBUG) println("\n\nIncoming packet:")
            if (DEBUG) hexdump(packet.data, 0, packet.getLength())
            val result = decode(packet.data, packet.getLength())
            callback?.onResult(result)
        }
    }

    /**
     * Ask for the A, SRV and TXT records of a particular service.
     *
     * @param serviceName the name of service to query in mDNS, e.g.
     * `"device-1234._example._tcp.local"`
     * @param timeout     duration in milliseconds to wait for an answer packet. If `0`, this
     * method will listen forever.
     * @return the reply packet's decoded answer data
     * @throws IOException
     */
    @Throws(IOException::class)
    fun resolve(serviceName: String, timeout: Int): Result {
        require(timeout >= 0)
        val group = InetAddress.getByName(MULTICAST_GROUP_ADDRESS)
        val sock = MulticastSocket() // binds to a random free source port
        if (DEBUG) println("Source port is " + sock.getLocalPort())
        if (DEBUG) println("Query packet:")
        val data = queryPacket(
            serviceName,
            QCLASS_INTERNET.toInt() or CLASS_FLAG_UNICAST.toInt(),
            QTYPE_A.toInt(),
            QTYPE_SRV.toInt(),
            QTYPE_TXT.toInt()
        )
        if (DEBUG) hexdump(data, 0, data.size)
        var packet = DatagramPacket(data, data.size, group, PORT)
        sock.timeToLive = 255
        sock.send(packet)
        val buf = ByteArray(1024)
        packet = DatagramPacket(buf, buf.size)
        val result = Result()
        var endTime: Long = 0
        if (timeout != 0) {
            endTime = System.currentTimeMillis() + timeout
        }
        // records could be returned in different packets, so we have to loop
        // timeout applies to the acquisition of ALL packets
        while (result.a == null || result.srv == null || result.txt == null) {
            if (timeout != 0) {
                val remaining = (endTime - System.currentTimeMillis()).toInt()
                if (remaining <= 0) {
                    break
                }
                sock.soTimeout = remaining
            }
            sock.receive(packet)
            if (DEBUG) println("\n\nIncoming packet:")
            if (DEBUG) hexdump(packet.data, 0, packet.getLength())
            decode(packet.data, packet.getLength(), result)
        }
        return result
    }

    @Throws(IOException::class)
    private fun writeFQDN(name: String, out: OutputStream) {
        for (part in name.split("\\.".toRegex()).dropLastWhile { it.isEmpty() }.toTypedArray()) {
            out.write(part.length)
            out.write(part.toByteArray())
        }
        out.write(0)
    }

    private fun hexdump(data: ByteArray, offset: Int, length: Int) {
        var offset = offset
        while (offset < length) {
            System.out.printf("%08x", offset)
            val origOffset = offset
            var col: Int = 0
            while (col < 16 && offset < length) {
                System.out.printf(" %02x", data[offset].toInt() and 0xFF)
                col++
                offset++
            }
            while (col < 16) {
                System.out.printf("   ")
                col++
            }
            print(" ")
            offset = origOffset
            col = 0
            while (col < 16 && offset < length) {
                val `val` = data[offset]
                val c: Char
                if (`val` >= 32 && `val` < 127) {
                    c = Char(`val`.toUShort())
                } else {
                    c = '.'
                }
                System.out.printf("%c", c)
                col++
                offset++
            }
            println()
        }
    }

    @Throws(IOException::class)
    fun decode(packet: ByteArray, packetLength: Int): Result {
        val result = Result()
        decode(packet, packetLength, result)
        return result
    }

    @Throws(IOException::class)
    fun decode(packet: ByteArray, packetLength: Int, result: Result) {
        val dis = DataInputStream(ByteArrayInputStream(packet, 0, packetLength))
        val transactionID = dis.readShort()
        val flags = dis.readShort()
        val questions = dis.readUnsignedShort()
        val answers = dis.readUnsignedShort()
        val authorityRRs = dis.readUnsignedShort()
        val additionalRRs = dis.readUnsignedShort()
        // decode the queries
        for (i in 0..<questions) {
            val fqdn = decodeFQDN(dis, packet, packetLength)
            val type = dis.readShort()
            val qclass = dis.readShort()
        }
        // decode the answers
        for (i in 0..<answers + authorityRRs + additionalRRs) {
            val fqdn = decodeFQDN(dis, packet, packetLength)
            val type = dis.readShort()
            val aclass = dis.readShort()
            if (DEBUG) System.out.printf("%s record%n", typeString(type))
            if (DEBUG) println("Name: " + fqdn)
            val ttl = dis.readInt()
            val length = dis.readUnsignedShort()
            val data = ByteArray(length)
            dis.readFully(data)
            var record: Record? = null
            when (type) {
                QTYPE_A -> {
                    result.a = decodeA(data)
                    record = result.a
                }

                QTYPE_SRV -> {
                    result.srv = decodeSRV(data, packet, packetLength)
                    record = result.srv
                }

                QTYPE_PTR -> decodePTR(data, packet, packetLength)
                QTYPE_TXT -> {
                    result.txt = decodeTXT(data)
                    record = result.txt
                }

                else -> if (DEBUG) hexdump(data, 0, data.size)
            }
            if (record != null) {
                record.fqdn = fqdn
                record.ttl = ttl
            }
        }
    }

    @Throws(IOException::class)
    private fun decodeSRV(srvData: ByteArray, packetData: ByteArray, packetLength: Int): SRV {
        val dis = DataInputStream(ByteArrayInputStream(srvData))
        val srv = SRV()
        srv.priority = dis.readUnsignedShort()
        srv.weight = dis.readUnsignedShort()
        srv.port = dis.readUnsignedShort()
        srv.target = decodeFQDN(dis, packetData, packetLength)
        if (DEBUG) System.out.printf(
            "Priority: %d Weight: %d Port: %d Target: %s%n",
            srv.priority,
            srv.weight,
            srv.port,
            srv.target
        )
        return srv
    }

    private fun typeString(type: Short): String {
        return when (type) {
            QTYPE_A -> "A"
            QTYPE_PTR -> "PTR"
            QTYPE_SRV -> "SRV"
            QTYPE_TXT -> "TXT"
            else -> "Unknown"
        }
    }

    @Throws(IOException::class)
    private fun decodePTR(ptrData: ByteArray, packet: ByteArray, packetLength: Int): String {
        val dis = DataInputStream(ByteArrayInputStream(ptrData))
        val fqdn = decodeFQDN(dis, packet, packetLength)
        if (DEBUG) println(fqdn)
        return fqdn
    }

    @Throws(IOException::class)
    private fun decodeA(data: ByteArray): A {
        if (data.size < 4) throw IOException("expected 4 bytes for IPv4 addr")
        val a = A()
        a.ipaddr =
            (data[0].toInt() and 0xFF).toString() + "." + (data[1].toInt() and 0xFF) + "." + (data[2].toInt() and 0xFF) + "." + (data[3].toInt() and 0xFF)
        if (DEBUG) println("Ipaddr: " + a.ipaddr)
        return a
    }

    @Throws(IOException::class)
    private fun decodeTXT(data: ByteArray): TXT {
        val txt = TXT()
        txt.dict = HashMap<String, String>()
        val dis = DataInputStream(ByteArrayInputStream(data))
        while (true) {
            val length: Int
            try {
                length = dis.readUnsignedByte()
            } catch (e: EOFException) {
                return txt
            }
            val segmentBytes = ByteArray(length)
            dis.readFully(segmentBytes)
            val segment = String(segmentBytes)
            val pos = segment.indexOf('=')
            val key: String?
            var value: String? = null
            if (pos != -1) {
                key = segment.substring(0, pos)
                value = segment.substring(pos + 1)
            } else {
                key = segment
            }
            if (DEBUG) println("$key=$value")
            if (!txt.dict!!.containsKey(key)) {
                // from RFC6763
                // If a client receives a TXT record containing the same key more than once, then
                // the client MUST silently ignore all but the first occurrence of that attribute."
                txt.dict!!.put(key, value?:"")
            }
        }
    }

    @Throws(IOException::class)
    private fun decodeFQDN(dis: DataInputStream, packet: ByteArray, packetLength: Int): String {
        var dis = dis
        val result = StringBuilder()
        var dot = false
        while (true) {
            var pointerHopCount = 0
            var length: Int
            while (true) {
                length = dis.readUnsignedByte()
                if (length == 0) return result.toString()
                if ((length and 0xc0) == 0xc0) {
                    // this is a compression method, the remainder of the string is a pointer to elsewhere in the packet
                    // adjust the stream boundary and repeat processing
                    if ((++pointerHopCount) * 2 >= packetLength) {
                        // We must have visited one of the possible pointers more than once => cycle
                        // this doesn't add to the domain length, but decoding would be non-terminating
                        throw IOException("cyclic empty references in domain name")
                    }
                    length = length and 0x3f
                    val offset = (length shl 8) or dis.readUnsignedByte()
                    dis = DataInputStream(ByteArrayInputStream(packet, offset, packetLength - offset))
                } else {
                    break
                }
            }
            val segment = ByteArray(length)
            dis.readFully(segment)
            if (dot) result.append('.')
            dot = true
            result.append(String(segment))
            if (result.length > packetLength) {
                // If we get here, we must be following cyclic references, since non-cyclic
                // references can't encode a domain name longer than the total length of the packet.
                // The domain name would be infinitely long, so abort now rather than consume
                // maximum heap.
                throw IOException("cyclic non-empty references in domain name")
            }
        }
    }

    /**
     * @see .discover
     */
    interface Callback {
        fun onResult(result: Result)
    }

    open class Record {
        /**
         * Fully-Qualified Domain Name of the record.
         */
        var fqdn: String? = null

        /**
         * Time-to-live of the record, in seconds.
         */
        var ttl: Int = 0
    }

    /**
     * DNS A record
     */
    class A : Record() {
        /**
         * The IPv4 address in dot-decimal notation, e.g. `"192.168.1.100"`
         */
        var ipaddr: String? = null
    }

    class SRV : Record() {
        var priority: Int = 0
        var weight: Int = 0
        var port: Int = 0

        /**
         * Fully-Qualified Domain Name of the target service.
         */
        var target: String? = null
    }

    class TXT : Record() {
        /**
         * The content of the TXT record's key-value store decoded as a [Map]
         */
        var dict: MutableMap<String, String>? = null
    }

    /**
     * Represents the decoded content of the answer sections of an incoming packet.
     * When the corresponding data is present in an answer, fields will be initialized with
     * populated data structures. When no such answer is present in the packet, fields will be
     * `null`.
     */
    class Result {
        var a: A? = null
        var srv: SRV? = null
        var txt: TXT? = null
    }
}
